import torch

def get_vector_field_component(u, dim):
    # 提取矢量场u中的第dim个分量,dim=0表示x分量，dim=1表示y分量，dim=2表示z分量
    # 矢量场u的维度分别是：[..., 3]
    # 输出的维度是：[..., 1]
    return u[..., dim].unsqueeze(-1)

def auto_grad(u, xyz_field, dim, requires_grad=True):
    """
    计算矢量场 u 的梯度（ux, uy 或 uz）。

    参数:
    u (torch.Tensor): 输入矢量场，维度为 [..., 3]。
    xyz_field (torch.Tensor): 坐标场，维度为 [..., 3]。
    dim (int): 要计算梯度的分量索引，0表示x分量，1表示y分量，2表示z分量。
    requires_grad (bool): 是否需要计算高阶梯度（用于二阶导数等）。

    返回:
    torch.Tensor: 梯度场，维度与 u 相同。
    """

    # 提取指定分量
    u_xyz = u[..., dim].unsqueeze(-1)  # 直接提取分量并增加维度

    # 创建与 u_xyz 相同形状的全1张量，放在GPU上
    ones = torch.ones_like(u_xyz, device='cuda')

    # 计算梯度
    if requires_grad:
        grad_u_xyz = torch.autograd.grad(
            outputs=u_xyz,
            inputs=xyz_field,
            grad_outputs=ones,
            create_graph=True,
            retain_graph=True,
            allow_unused=False
        )[0]
    else:
        with torch.no_grad():  # 禁用梯度计算以提升速度
            grad_u_xyz = torch.autograd.grad(
                outputs=u_xyz,
                inputs=xyz_field,
                grad_outputs=ones,
                create_graph=False,
                retain_graph=True,
                allow_unused=False
            )[0]

    return grad_u_xyz


if __name__ == '__main__':
    # 定义一个简单的函数
    def f(x):
        return x**2

    # 创建一个张量,模拟坐标矢量场
    # 矢量场x的维度分别是：[batch_size, x, y, z, 3]
    x = torch.tensor([[[[[2.0, 3.0, 5.0],
                         [-2.0, -3.0, -5.0]]]]], requires_grad=True, device='cuda')
    print("x:", x)

    # 计算函数的值（模拟位移）
    u = f(x)
    print("u:", u)

    # 求导，模拟位移梯度
    # ————————————————————以下是典型用例————————————————————————
    duxdxyz = auto_grad(u, x, 0)
    duydxyz = auto_grad(u, x, 1)
    duzdxyz = auto_grad(u, x, 2)
    # ————————————————————以上是典型用例————————————————————————

    # 提取分量，模拟变形梯度场
    Fxx = get_vector_field_component(duxdxyz, 0) + 1
    Fxy = get_vector_field_component(duxdxyz, 1) + 0
    Fxz = get_vector_field_component(duxdxyz, 2) + 0
    Fyx = get_vector_field_component(duydxyz, 0) + 0
    Fyy = get_vector_field_component(duydxyz, 1) + 1
    Fyz = get_vector_field_component(duydxyz, 2) + 0
    Fzx = get_vector_field_component(duzdxyz, 0) + 0
    Fzy = get_vector_field_component(duzdxyz, 1) + 0
    Fzz = get_vector_field_component(duzdxyz, 2) + 1

    print("dux/dx, dux/dy, dux/dz:", duxdxyz)
    print("duy/dx, duy/dy, duy/dz:", duydxyz)
    print("duz/dx, duz/dy, duz/dz:", duzdxyz)
    print("Fxx", Fxx)
    print("Fxy", Fxy)
    print("Fxz", Fxz)
